######################
###### Plotting#######
#####################
data_anemia <- read_dta("Jehanabad_analytic_sample_hemoglobin.dta")
# read in data
data_cog <- read_dta("Jehanabad_analytic_sample_cognition.dta")

data_hemo <- data_anemia %>% mutate(
    any_anemia = ifelse(hemoglobin < 11.5, 1, 0),
    mild_anemia = ifelse(hemoglobin >= 11 & hemoglobin < 11.5, 1, 0),
    moderate_severe_anemia = ifelse(hemoglobin < 11, 1, 0)
)

######## function treatment ############
### takes in a couple of arguments and returns a plot and saves it automatically
# with automatic smoothing for local regression fitting, smoothing is by 
# default high
# plots treatment effect vs attendance
# @dataset to base analysis on
# @dep_var dependent variable
# @y_axis label to be specified in axis title
# @points whether points should be included
# @limit whether we want to truncate the attendance rate,eg. to 70%


treatment <- function(dataset, dep_var, y_axis, points = T, limit = NA) {
  formula <- paste(paste(dep_var, "~ "),
    paste(c(
      "no_HHmembers", "years_schooling_father",
      "years_schooling_mother",
      "asset_index", "total_enrollment", "class_size",
      "student_teacher_ratio",
      "wave:treat:factor(total_attendance_2)",
      "treat*wave",
      "treat", "wave"
    ), collapse = "+"),
    sep = ""
  )
  df <- lm(formula,
    data = dataset
  )
  row_names <- rownames(data.frame(df$coefficients))
  index <- 12:length(df$coefficients)
  y <- as.numeric(df$coefficients)[11] + as.numeric(df$coefficients)[index]
  x <- str_remove(
    as.character(row_names),
    "wave\\:treat\\:factor\\(total_attendance_2\\)"
  )[index]

  plot_input <- data.frame(y = y, x = as.character(x))

  plot_input <- plot_input %>%
    mutate(x = as.numeric(as.character(x)), y = as.numeric(as.character(y)))
  plot <- ggplot(plot_input, aes(x, y)) +
    geom_smooth(method = "loess") +
    theme(axis.text.x = element_text(angle = 90, hjust = 1))
  if (points == T) {
    plot <- plot + geom_point()
  } else {
    plot <- plot + geom_point(alpha = 0)
  }
  plot <- plot +
    ylab(paste("DID treatment effect on \n", y_axis)) +
    xlab("Child's school attendance") +
    xlim(0.5, 1) +
    ggtitle("") +
    theme(
      axis.text = element_text(size = 7),
      axis.title.y = element_text(size = 7, vjust = -0.65)
    )
  if (!is.na(limit)) {
    if (points == T) {
      ending <- paste(limit, "points.png")
    } else {
      ending <- paste(limit, ".png")
    }
  } else {
    if (points == T) {
      ending <- "points01.png"
    } else {
      ending <- "01.png"
    }
  }

  ggsave(paste(y_axis, ending))
  return(plot)
}

treatment(data_cog, "day_and_night_norm",
  y_axis = "day and night", points = T,
  limit = 0.5
)
treatment(data_hemo, "hemoglobin",
  y_axis = "hemoglobin", points = T,
  limit = 0.5
)
treatment(data_cog, "Cognition_Score_1_norm",
  y_axis = "cognition",
  points = T, limit = 0.5
)
treatment(data_cog, "reading_norm",
  y_axis = "reading", points = T,
  limit = 0.5
)
treatment(data_hemo, "math_norm", y_axis = "math", points = T, limit = 0.5)
treatment(data_hemo, "any_anemia",
  y_axis = "any anemia", points = T,
  limit = 0.5
)
